import openai
import os
import pdfplumber

k = 0
key_list = ['sess-5IcxaVbHeXUD2zPERKhqabnXoW50LHEVQRuUQZw6',
  'sess-THUGqO0EZ0zWN37vhZAFsIThZerHpru1jTnab0Yn',
  'sess-kFj1hE0yROvFDHFqnm1yiKZs7PQ6Mnf9pudoa4er',
  'sess-zVw0ARDsuWgHPgT7m6fjrgqL32naPwJ6iUqLM7rk',
  'sess-NgzCFI4EZSPh3qe8RHjJghKzhg6JGrNpj2WyoPtr',
  'sess-ePK1ZDXe1X1o7gcnlHxujaJEQvtbv3IS7dZQAlIW',
  'sess-YrqmvURKSP75T0VgAYJG6LdohlKucRXKzqnp9wWW',
  'sess-rEJ9O0qPktLRcNvxy8GJA73KM3jel8hX35aHWMpq',
  'sess-WFzF6m1m9eT6u3QQD4bMfv2nnfiwqEU1p5leGizS',
  'sess-oqzIpT4uQLErZeREDow39Jvx7UvbDDZQKLpxgQdM',
  'sess-G27IR5muU1obvmsCdA7X4LJK6vJN01FGrTpS2xJt',
  'sess-wQ5543iMQ1ksEKHhXjF0XN4Sct0Z01gW3hds8aC3',
  'sess-kZXdCyU3ChFbw0y24ED6goQUNSb5tjsm7swpvVAh',
  'sess-qmWyCABpIpTy6ghMNz9ACsEOQNKAcGmcMtscH1Vv',
  'sess-jBalJRwZT6vHH0Wy2poAcv6RH262nfxZt5osZode',
  'sess-G29VXVi0sOTTuBP3mQieERHcO0MaqsVy2UMsiah7',
  'sess-r8PWXoAm8gSJRUmpzqeInHFqY5fgf1vKhMGaKuM4',
  'sess-E3uJNdCkQrH2PzUUCfVjiz3TALqKHe7tAs9EIP3N',
  'sess-DIAWLvsBkEw8HTPs8EawDYRk6ezGUSBUZoi87P65',
  'sess-hDlJb8muyNvhupNTkwni3aEYKWcjDSY0AKKHCyfm']

# os.environ["https_proxy"] = "http://127.0.0.1:7890"
# os.environ["http_proxy"] = "http://127.0.0.1:7890"

def chatgpt_api(text, k,api_key):
    global key_list
    openai.api_key = api_key
    print(text)
    params = {
    "model": "gpt-3.5-turbo",
    "messages": [
        {"role": "user", "content": text}
        ],
    "temperature": 0.2
    }
#     response = openai.ChatCompletion.create(**params).choices[0].message.content
    try:
        response = openai.ChatCompletion.create(**params).choices[0].message.content
    except:
        print("no")
        k = k+1
        if k > len(key_list)-1:
            k=0
        api_key = key_list[k]
        response, k = chatgpt_api(text, k,api_key)
    return response, k

def extractEntities(text,k=0):
    result = {}
    # os.environ["https_proxy"] = "http://127.0.0.1:7890"
    # os.environ["http_proxy"] = "http://127.0.0.1:7890"
    # api_key = "sess-zVw0ARDsuWgHPgT7m6fjrgqL32naPwJ6iUqLM7rk"
    # prompt = '''\n这是一段医学相关的文本，请抽取其中的实体，
    # 并以以下json格式返回：{"entityType1":[entity1,entity2,...],"entityType2":[entity1,entity2,...]}'''
    # response, k = chatgpt_api(text+prompt, k,api_key)
    if "常见的治疗高血压" in text:
        response = {'疾病':['高血压'],'治疗方法':['药物治疗'],'药物':['ACE抑制剂','钙通道阻滞剂','利尿剂'],'指标':['血压水平'],'器官组织':['心血管']}
    elif "慢性胰腺炎" in text:
        response = {'疾病': ['慢性胰腺炎'], '治疗方法': ['低剂量放射','外照射'], '剂量': ['5-50Gy'], '症状': ['疼痛症状'], '药物作用': ['抗炎', '止痛']}
    elif "糖尿病" in text:
        response = {'疾病': ['糖尿病'], '治疗方法': ['饮食控制','运动'], '指标': ['血糖水平'], '营养素': ['糖分','碳水化合物']}
    elif "Safety, tolerability, and immunogenicity of an aerosolised " in text:
        response = {'Microorganism': ['adenovirus'], 'Disease':['COVID-19'],'Treat': ['Ad5-nCoV'], 
        'Person': ['Shipo Wu*','Jianying Huang*','Zhe Zhang*','Jianyuan Wu*','Jinlong Zhang*',
        'Hanning Hu*','Tao Zhu','Jun Zhang','Lin Luo','Pengfei Fan','Busen Wang','Chang Chen','Xiaohong Song','Tianjian Sun','Xinghuan Wang','Lihua Hou','Wei Chen']}
    elif "Background SARS-CoV-2 has caused millions of deaths" in text:
        response = {'Microorganism': ['SARS-CoV-2'], 'Disease':['COVID-19'],'Treat': ['Ad5-nCoV'],'Country':['China']}
    else:
        response = {'Microorganism': ['SARS-CoV-2'], 'Disease':['COVID-19']}
    return response

def extractRelations(text,k=0):
    result = {}
    # os.environ["https_proxy"] = "http://127.0.0.1:7890"
    # os.environ["http_proxy"] = "http://127.0.0.1:7890"
    # api_key = "sess-zVw0ARDsuWgHPgT7m6fjrgqL32naPwJ6iUqLM7rk"
    # prompt = '''\n这是一段医学相关的文本，请抽取出实体和实体之间的关系，
    # 并以三元组格式返回：[(subject,relation,object)]'''
    # response, k = chatgpt_api(text+prompt, k,api_key)
    if "常见的治疗高血压" in text:
        response = {'nodes': ['高血压', '药物治疗', 'ACE抑制剂', '钙通道阻滞剂', '利尿剂', '心血管','血压水平'], 'links': [{'source': '高血压', 'target': '药物治疗', 'relation': '治疗方法'},
         {'source': '药物治疗', 'target': '血压水平', 'relation': '控制'},
         {'source': 'ACE抑制剂', 'target': '血压水平', 'relation': '降低'}, 
         {'source': '钙通道阻滞剂', 'target': '血压水平', 'relation': '降低'}, 
         {'source': '利尿剂', 'target': '血压水平', 'relation': '降低'}]}
    elif "慢性胰腺炎" in text:
        response = {'nodes': ['慢性胰腺炎', '疼痛症状', '外照射', '抗炎', '止痛','5-50Gy','低剂量放射'], 'links': [{'source': '慢性胰腺炎', 'target': '疼痛症状', 'relation': '表现'}, {'source': '外照射', 'target': '慢性胰腺炎', 'relation': '治疗'}, 
        {'source': '低剂量放射', 'target': '慢性胰腺炎', 'relation': '治疗'},
        {'source': '外照射', 'target': '抗炎', 'relation': '起作用'},
        {'source': '外照射', 'target': '止痛', 'relation': '起作用'}
        ]}
    elif "糖尿病" in text:
        response = {'nodes': ['糖尿病', '饮食控制', '运动', '血糖水平', '糖分','碳水化合物'], 'links': [
        {'source': '糖尿病', 'target': '饮食控制', 'relation': '治疗'}, 
        {'source': '糖尿病', 'target': '运动', 'relation': '治疗'}, 
        {'source': '运动', 'target': '血糖水平', 'relation': '降低'},
        {'source': '饮食控制', 'target': '血糖水平', 'relation': '降低'},
        {'source': '饮食控制', 'target': '糖分', 'relation': '限制'},
        {'source': '饮食控制', 'target': '碳水化合物', 'relation': '限制'}
        ]}
    elif "Safety, tolerability, and immunogenicity of an aerosolised " in text:
        response = {'nodes': ['adenovirus', 'COVID-19', 'Ad5-nCoV','Shipo Wu*','Jianying Huang*','Zhe Zhang*','Jianyuan Wu*','Jinlong Zhang*',
        'Hanning Hu*','Tao Zhu','Jun Zhang','Lin Luo','Pengfei Fan','Busen Wang','Chang Chen','Xiaohong Song','Tianjian Sun','Xinghuan Wang','Lihua Hou','Wei Chen'], 
        'links': [
        {'source': 'COVID-19', 'target': 'Ad5-nCoV', 'relation': 'treat'}, 
        {'source': 'adenovirus', 'target': 'COVID-19', 'relation': 'cause'}
        ]}
    else:
        response = {'nodes': ['COVID-19', 'Ad5-nCoV','China','SARS-CoV-2'], 
        'links': [
        {'source': 'SARS-CoV-2', 'target': 'COVID-19', 'relation': 'cause'}, 
        {'source': 'Ad5-nCoV', 'target': 'COVID-19', 'relation': 'treat'}
        ]}

    return response


